'''
Author: 
Date: 2022-01-06 18:34:00
LastEditTime: 2022-12-08 15:10:59
LastEditors: Please set LastEditors

'''

import math
import matplotlib.pyplot as plt
import sys
import importlib
import os
import time
from loguru import logger
import cv2
import torch
from yolox.data.data_augment import ValTransform
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import get_model_info, postprocess

from pill_visualize import visual


class Predictor(object):
    def __init__(
        self,
        model,
        exp,
        cls_names=COCO_CLASSES,
        trt_file=None,
        decoder=None,
        device="gpu",
        fp16=False,
        legacy=False,
    ):
        self.model = model
        self.cls_names = cls_names
        self.decoder = decoder
        self.num_classes = exp.num_classes
        self.confthre = exp.test_conf
        self.nmsthre = exp.nmsthre
        self.test_size = exp.test_size
        self.device = device
        self.fp16 = fp16
        self.preproc = ValTransform(legacy=legacy)
        if trt_file is not None:
            from torch2trt import TRTModule
            model_trt = TRTModule()
            model_trt.load_state_dict(torch.load(trt_file))
            x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()
            self.model(x)
            self.model = model_trt

    def inference(self, img):
        img_info = {"id": 0}
        if isinstance(img, str):
            img_info["file_name"] = os.path.basename(img)
            img = cv2.imread(img)
        else:
            img_info["file_name"] = None

        height, width = img.shape[:2]
        img_info["height"] = height
        img_info["width"] = width
        img_info["raw_img"] = img

        ratio = min(self.test_size[0] / img.shape[0], self.test_size[1] / img.shape[1])
        img_info["ratio"] = ratio

        img, _ = self.preproc(img, None, self.test_size)
        img = torch.from_numpy(img).unsqueeze(0)
        img = img.float()
        if self.device == "gpu":
            img = img.cuda()
            if self.fp16:
                img = img.half()  # to FP16

        t0 = time.time()

        with torch.no_grad():            
            outputs = self.model(img)
            if self.decoder is not None:
                outputs = self.decoder(outputs, dtype=outputs.type())
            outputs = postprocess(
                outputs, self.num_classes, self.confthre,
                self.nmsthre, class_agnostic=True
            )
            #############################
            # logger.info("Infer time: {:.4f}s".format(time.time() - t0))
            # logger.info(len(outputs[0][:, 0:4]))

        img_info["infer_time"] = time.time() - t0

        return outputs, img_info


class YoloxHandler:
    def __init__(self, exp_file, ckpt_file=None, isTrt=False):
        try:
            sys.path.append(os.path.dirname(exp_file))
            current_exp = importlib.import_module(os.path.basename(exp_file).split(".")[0])
            exp = current_exp.Exp()
        except Exception:
            raise ImportError("{} doesn't contains class named 'Exp'".format(exp_file))

        self.exp = exp

        ############################################################
        self.exp.test_conf = 0.5  # 置信度设置
        self.exp.nmsthre = 0.45   # 非极大值抑制设置
        self.exp.test_size = (640, 640)
        ##############################################################

        self.exp.device = 'gpu'

        model = self.exp.get_model()
        logger.info("Model Summary: {}".format(get_model_info(model, exp.test_size)))

        model.cuda()

        model.eval()

        assert ckpt_file is not None, 'ckpt file is None'
        assert os.path.exists(ckpt_file), 'ckpt file not exists'

        logger.info("loading checkpoint")

        if isTrt:
            trt_file = ckpt_file
            model.head.decode_in_inference = False
            decoder = model.head.decode_outputs
            logger.info("Using TensorRT to inference")
        else:
            trt_file = None
            decoder = None
            ckpt = torch.load(ckpt_file, map_location="cpu")

            # load the model state dict
            model.load_state_dict(ckpt["model"])

        logger.info("loaded checkpoint done.")

        self.predictor = Predictor(
            model, exp, COCO_CLASSES, trt_file, decoder,
        )

    def predict(self, image):
        assert image is not None, 'image path is None'
        return self.predictor.inference(image)

    def visual(self, output, img_info, cls_conf=0.35):
        return visual(output, img_info, cls_conf,COCO_CLASSES)


if __name__ == "__main__": #这个是来临时测试的
    # 模型
    exp_file = "d:/ai/yolox/exps/example/custom/yolox_s.py"

    # 训练好的参数
    ckpt = "D:/AI/YOLOX/YOLOX_outputs/yolox_s/last_epoch_ckpt.pth"
    handle = YoloxHandler(exp_file, ckpt, False)

    # 预测
    fig = plt.figure(figsize=(6, 7))

    images = [
        "D:/AI/YOLOX/datasets/coco/val2017/001261.jpg",        
        "D:/AI/YOLOX/datasets/coco/val2017/000707.jpg",
        "D:/AI/YOLOX/datasets/coco/val2017/001221.jpg",
    ]

    i = 1
    xl = math.ceil(len(images)/2.0)  # 两行显示

    for img in images:

        if not os.path.exists(img):
            logger.warning('image not exist: %s' % img)
            continue

        output, img_info = handle.predict(img)
        result_image = handle.visual(output[0], img_info, 0.5)

        img_rgb = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB)
        fig.add_subplot(2, xl, i)
        plt.axis('off')
        plt.imshow(img_rgb)
        i += 1
    plt.subplots_adjust(wspace=0, hspace=0)  # 图片间隔
    plt.show()
